// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <melon/container/f14_set.h>

#include <re2/re2.h>
#include <pollux/common/base/spill_config.h>
#include <pollux/common/base/spill_stats.h>
#include <pollux/common/compression/compression.h>
#include <pollux/common/file/file.h>
#include <pollux/common/file/file_systems.h>
#include <pollux/exec/spill_file.h>
#include <pollux/exec/tree_of_losers.h>
#include <pollux/exec/unordered_stream_reader.h>
#include <pollux/vector/complex_vector.h>
#include <pollux/vector/decoded_vector.h>
#include <pollux/vector/vector_stream.h>

namespace kumo::pollux::exec {
    /// A source of sorted spilled RowVectors coming either from a file or memory.
    class SpillMergeStream : public MergeStream {
    public:
        SpillMergeStream() = default;

        virtual ~SpillMergeStream() = default;

        /// Returns the id of a spill merge stream which is unique in the merge set.
        virtual uint32_t id() const = 0;

        bool hasData() const final {
            return index_ < size_;
        }

        bool operator<(const MergeStream &other) const final {
            return compare(other) < 0;
        }

        int32_t compare(const MergeStream &other) const override;

        void pop();

        const RowVector &current() const {
            POLLUX_CHECK(!closed_);
            return *rowVector_;
        }

        /// Invoked to get the current row index in 'rowVector_'. If 'isLastRow' is
        /// not null, it is set to true if current row is the last one in the current
        /// batch, in which case the caller must call copy out current batch data if
        /// required before calling pop().
        vector_size_t currentIndex(bool *isLastRow = nullptr) const {
            POLLUX_CHECK(!closed_);
            if (isLastRow != nullptr) {
                *isLastRow = (index_ == (rowVector_->size() - 1));
            }
            return index_;
        }

        /// Returns a DecodedVector set decoding the 'index'th child of 'rowVector_'
        DecodedVector &decoded(int32_t index) {
            POLLUX_CHECK(!closed_);
            ensureDecodedValid(index);
            return decoded_[index];
        }

    protected:
        virtual int32_t numSortKeys() const = 0;

        virtual const std::vector<CompareFlags> &sortCompareFlags() const = 0;

        virtual void nextBatch() = 0;

        virtual void close();

        // loads the next 'row_vector' and sets 'decoded_' if this is initialized.
        void setNextBatch() {
            nextBatch();
            if (!decoded_.empty()) {
                ensureRows();
                for (auto i = 0; i < decoded_.size(); ++i) {
                    decoded_[i].decode(*rowVector_->childAt(i), rows_);
                }
            }
        }

        void ensureDecodedValid(int32_t index) {
            int32_t oldSize = decoded_.size();
            if (index < oldSize) {
                return;
            }
            ensureRows();
            decoded_.resize(index + 1);
            for (auto i = oldSize; i <= index; ++i) {
                decoded_[index].decode(*rowVector_->childAt(index), rows_);
            }
        }

        void ensureRows() {
            if (rows_.size() != rowVector_->size()) {
                rows_.resize(size_);
            }
        }

        // True if the stream is closed.
        bool closed_{false};

        // Current batch of rows.
        RowVectorPtr rowVector_;

        // The current row in 'rowVector_'
        vector_size_t index_{0};

        // Number of rows in 'rowVector_'
        vector_size_t size_{0};

        // Decoded vectors for leading parts of 'rowVector_'. Initialized on first
        // use and maintained when updating 'rowVector_'.
        std::vector<DecodedVector> decoded_;

        // Covers all rows inn 'rowVector_' Set if 'decoded_' is non-empty.
        SelectivityVector rows_;
    };

    /// A source of spilled RowVectors coming from a file.
    class FileSpillMergeStream : public SpillMergeStream {
    public:
        static std::unique_ptr<SpillMergeStream> create(
            std::unique_ptr<SpillReadFile> spillFile) {
            auto spillStream = std::unique_ptr<SpillMergeStream>(
                new FileSpillMergeStream(std::move(spillFile)));
            static_cast<FileSpillMergeStream *>(spillStream.get())->nextBatch();
            return spillStream;
        }

        uint32_t id() const override;

    private:
        explicit FileSpillMergeStream(std::unique_ptr<SpillReadFile> spillFile)
            : spillFile_(std::move(spillFile)) {
            POLLUX_CHECK_NOT_NULL(spillFile_);
        }

        int32_t numSortKeys() const override {
            POLLUX_CHECK(!closed_);
            return spillFile_->numSortKeys();
        }

        const std::vector<CompareFlags> &sortCompareFlags() const override {
            POLLUX_CHECK(!closed_);
            return spillFile_->sortCompareFlags();
        }

        void nextBatch() override;

        void close() override;

        std::unique_ptr<SpillReadFile> spillFile_;
    };

    /// A source of spilled RowVectors coming from a file. The spill data might not
    /// be sorted.
    ///
    /// NOTE: this object is not thread-safe.
    class FileSpillBatchStream : public BatchStream {
    public:
        static std::unique_ptr<BatchStream> create(
            std::unique_ptr<SpillReadFile> spillFile) {
            auto *spillStream = new FileSpillBatchStream(std::move(spillFile));
            return std::unique_ptr<BatchStream>(spillStream);
        }

        bool nextBatch(RowVectorPtr &batch) override {
            return spillFile_->nextBatch(batch);
        }

    private:
        explicit FileSpillBatchStream(std::unique_ptr<SpillReadFile> spillFile)
            : spillFile_(std::move(spillFile)) {
            POLLUX_CHECK_NOT_NULL(spillFile_);
        }

        std::unique_ptr<SpillReadFile> spillFile_;
    };

    /// Identifies a spill partition generated from a given spilling operator. It
    /// consists of partition start bit offset and the actual partition number.
    /// The start bit offset is used to calculate the partition number of spill
    /// data. It is required for the recursive spilling handling as we advance the
    /// start bit offset when we go to the next level of recursive spilling.
    ///
    /// NOTE: multiple shards created from the same SpillPartition by split()
    /// will share the same id.
    class SpillPartitionId {
    public:
        SpillPartitionId(uint8_t partitionBitOffset, int32_t partitionNumber)
            : partitionBitOffset_(partitionBitOffset),
              partitionNumber_(partitionNumber) {
        }

        bool operator==(const SpillPartitionId &other) const {
            return std::tie(partitionBitOffset_, partitionNumber_) ==
                   std::tie(other.partitionBitOffset_, other.partitionNumber_);
        }

        bool operator!=(const SpillPartitionId &other) const {
            return !(*this == other);
        }

        /// Customize the compare operator for recursive spilling control. It
        /// ensures the partition with higher partition bit is handled prior than
        /// one with lower partition bit. With the same partition bit, the one with
        /// smaller partition number is handled first. We put all spill partitions
        /// in an ordered map sorted based on the partition id. The recursive
        /// spilling will advance the partition start bit when go to the next level
        /// of recursive spilling.
        bool operator<(const SpillPartitionId &other) const {
            if (partitionBitOffset_ != other.partitionBitOffset_) {
                return partitionBitOffset_ > other.partitionBitOffset_;
            }
            return partitionNumber_ < other.partitionNumber_;
        }

        bool operator>(const SpillPartitionId &other) const {
            return (*this != other) && !(*this < other);
        }

        std::string toString() const {
            return fmt::format("[{},{}]", partitionBitOffset_, partitionNumber_);
        }

        uint8_t partitionBitOffset() const {
            return partitionBitOffset_;
        }

        int32_t partitionNumber() const {
            return partitionNumber_;
        }

    private:
        uint8_t partitionBitOffset_{0};
        int32_t partitionNumber_{0};
    };

    inline std::ostream &operator<<(std::ostream &os, SpillPartitionId id) {
        os << id.toString();
        return os;
    }

    using SpillPartitionIdSet = melon::F14FastSet<SpillPartitionId>;
    using SpillPartitionNumSet = melon::F14FastSet<uint32_t>;

    /// Contains a spill partition data which includes the partition id and
    /// corresponding spill files.
    class SpillPartition {
    public:
        explicit SpillPartition(const SpillPartitionId &id)
            : SpillPartition(id, {}) {
        }

        SpillPartition(const SpillPartitionId &id, SpillFiles files) : id_(id) {
            addFiles(std::move(files));
        }

        void addFiles(SpillFiles files) {
            files_.reserve(files_.size() + files.size());
            for (auto &file: files) {
                size_ += file.size;
                files_.push_back(std::move(file));
            }
        }

        const SpillPartitionId &id() const {
            return id_;
        }

        int numFiles() const {
            return files_.size();
        }

        /// Returns the total file byte size of this spilled partition.
        uint64_t size() const {
            return size_;
        }

        /// Invoked to split this spill partition into 'numShards' to process in
        /// parallel.
        ///
        /// NOTE: the split spill partition shards will have the same id as this.
        std::vector<std::unique_ptr<SpillPartition> > split(int numShards);

        /// Invoked to create an unordered stream reader from this spill partition.
        /// The created reader will take the ownership of the spill files.
        /// 'bufferSize' specifies the read size from the storage. If the file
        /// system supports async read mode, then reader allocates two buffers with
        /// one buffer prefetch ahead. 'spillStats' is provided to collect the spill
        /// stats when reading data from spilled files.
        std::unique_ptr<UnorderedStreamReader<BatchStream> > createUnorderedReader(
            uint64_t bufferSize,
            memory::MemoryPool *pool,
            melon::Synchronized<common::SpillStats> *spillStats);

        /// Invoked to create an ordered stream reader from this spill partition.
        /// The created reader will take the ownership of the spill files.
        /// 'bufferSize' specifies the read size from the storage. If the file
        /// system supports async read mode, then reader allocates two buffers with
        /// one buffer prefetch ahead. 'spillStats' is provided to collect the spill
        /// stats when reading data from spilled files.
        std::unique_ptr<TreeOfLosers<SpillMergeStream> > createOrderedReader(
            uint64_t bufferSize,
            memory::MemoryPool *pool,
            melon::Synchronized<common::SpillStats> *spillStats);

        std::string toString() const;

    private:
        SpillPartitionId id_;
        SpillFiles files_;
        // Counts the total file size in bytes from this spilled partition.
        uint64_t size_{0};
    };

    using SpillPartitionSet =
    std::map<SpillPartitionId, std::unique_ptr<SpillPartition> >;

    /// Represents all spilled data of an operator, e.g. order by or group
    /// by. This has one SpillFileList per partition of spill data.
    class SpillState {
    public:
        /// Constructs a SpillState. 'type' is the content RowType. 'path' is the
        /// file system path prefix. 'bits' is the hash bit field for partitioning
        /// data between files. This also gives the maximum number of partitions.
        /// 'numSortKeys' is the number of leading columns on which the data is
        /// sorted, 0 if only hash partitioning is used. 'targetFileSize' is the
        /// target size of a single file.  'pool' owns the memory for state and
        /// results.
        SpillState(
            const common::GetSpillDirectoryPathCB &getSpillDirectoryPath,
            const common::UpdateAndCheckSpillLimitCB &updateAndCheckSpillLimitCb,
            const std::string &fileNamePrefix,
            int32_t maxPartitions,
            int32_t numSortKeys,
            const std::vector<CompareFlags> &sortCompareFlags,
            uint64_t targetFileSize,
            uint64_t writeBufferSize,
            common::CompressionKind compressionKind,
            const std::optional<common::PrefixSortConfig> &prefixSortConfig,
            memory::MemoryPool *pool,
            melon::Synchronized<common::SpillStats> *stats,
            const std::string &fileCreateConfig = {});

        /// Indicates if a given 'partition' has been spilled or not.
        bool isPartitionSpilled(uint32_t partition) const {
            POLLUX_DCHECK_LT(partition, maxPartitions_);
            return spilledPartitionSet_.contains(partition);
        }

        // Sets a partition as spilled.
        void setPartitionSpilled(uint32_t partition);

        // Returns how many ways spilled data can be partitioned.
        int32_t maxPartitions() const {
            return maxPartitions_;
        }

        uint64_t targetFileSize() const {
            return targetFileSize_;
        }

        common::CompressionKind compressionKind() const {
            return compressionKind_;
        }

        const std::optional<common::PrefixSortConfig> &prefixSortConfig() const {
            return prefixSortConfig_;
        }

        const std::vector<CompareFlags> &sortCompareFlags() const {
            return sortCompareFlags_;
        }

        bool isAnyPartitionSpilled() const {
            return !spilledPartitionSet_.empty();
        }

        bool isAllPartitionSpilled() const {
            POLLUX_CHECK_LE(spilledPartitionSet_.size(), maxPartitions_);
            return spilledPartitionSet_.size() == maxPartitions_;
        }

        /// Appends data to 'partition'. The rows given by 'indices' must be sorted
        /// for a sorted spill and must hash to 'partition'. It is safe to call this
        /// on multiple threads if all threads specify a different partition.
        /// Returns the size to append to partition.
        uint64_t appendToPartition(uint32_t partition, const RowVectorPtr &rows);

        /// Finishes a sorted run for 'partition'. If write is called for
        /// 'partition' again, the data does not have to be sorted relative to the
        /// data written so far.
        void finishFile(uint32_t partition);

        /// Returns the current number of finished files from a given partition.
        ///
        /// NOTE: the fucntion returns zero if the state has finished or the
        /// partition is not spilled yet.
        size_t numFinishedFiles(uint32_t partition) const;

        /// Returns the spill file objects from a given 'partition'. The function
        /// returns an empty list if either the partition has not been spilled or
        /// has no spilled data.
        SpillFiles finish(uint32_t partition);

        /// Returns the spilled partition number set.
        const SpillPartitionNumSet &spilledPartitionSet() const;

        /// Returns the spilled file paths from all the partitions.
        std::vector<std::string> testingSpilledFilePaths() const;

        /// Returns the file ids from a given partition.
        std::vector<uint32_t> testingSpilledFileIds(int32_t partitionNum) const;

        /// Returns the set of partitions that have spilled data.
        SpillPartitionNumSet testingNonEmptySpilledPartitionSet() const;

    private:
        // Ensures that the bytes to spill is within the limit of
        // maxSpillBytesPerWrite_ for a given spill write/appendToPartition call.
        // NOTE: the Presto serializer used for spill serialization can't handle more
        // than 2GB data size. Hence we can't spill a vector which exceeds
        // 2GB serialized buffer.
        void validateSpillBytesSize(uint64_t bytes);

        void updateSpilledInputBytes(uint64_t bytes);

        SpillWriter *partitionWriter(uint32_t partition) const;

        const RowTypePtr type_;

        // A callback function that returns the spill directory path.
        // Implementations can use it to ensure the path exists before returning.
        common::GetSpillDirectoryPathCB getSpillDirPathCb_;

        // Updates the aggregated spill bytes of this query, and throws if exceeds
        // the max spill bytes limit.
        common::UpdateAndCheckSpillLimitCB updateAndCheckSpillLimitCb_;

        // Prefix for spill files.
        const std::string fileNamePrefix_;
        const int32_t maxPartitions_;
        const int32_t numSortKeys_;
        const std::vector<CompareFlags> sortCompareFlags_;
        const uint64_t targetFileSize_;
        const uint64_t writeBufferSize_;
        const common::CompressionKind compressionKind_;
        const std::optional<common::PrefixSortConfig> prefixSortConfig_;
        const std::string fileCreateConfig_;
        memory::MemoryPool *const pool_;
        melon::Synchronized<common::SpillStats> *const stats_;

        // A set of spilled partition numbers.
        SpillPartitionNumSet spilledPartitionSet_;

        // A file list for each spilled partition. Only partitions that have
        // started spilling have an entry here.
        std::vector<std::unique_ptr<SpillWriter> > partitionWriters_;
    };

    /// Generate partition id set from given spill partition set.
    SpillPartitionIdSet toSpillPartitionIdSet(
        const SpillPartitionSet &partitionSet);

    /// Scoped spill percentage utility that allows user to set the behavior of
    /// triggered spill.
    /// 'spillPct' indicates the chance of triggering spilling. 100% means spill
    /// will always be triggered.
    /// 'pools' is a regular expression string used to define the specific memory
    /// pools targeted for injecting spilling.
    /// 'maxInjections' indicates the max number of actual triggering. e.g. when
    /// 'spillPct' is 20 and 'maxInjections' is 10, continuous calls to
    /// testingTriggerSpill(poolName) will keep rolling the dice that has a
    /// chance of 20% triggering until 10 triggers have been invoked.
    class TestScopedSpillInjection {
    public:
        explicit TestScopedSpillInjection(
            int32_t spillPct,
            const std::string &poolRegExp = ".*",
            uint32_t maxInjections = std::numeric_limits<uint32_t>::max());

        ~TestScopedSpillInjection();
    };

    /// Test utility that returns true if triggered spill is evaluated to happen,
    /// false otherwise.
    bool testingTriggerSpill(const std::string &poolName = "");

    tsan_atomic<uint32_t> &injectedSpillCount();

    /// Removes empty partitions from given spill partition set.
    void removeEmptyPartitions(SpillPartitionSet &partitionSet);
} // namespace kumo::pollux::exec

// Adding the custom hash for SpillPartitionId to std::hash to make it usable
// with maps and other standard data structures.
namespace std {
    template<>
    struct hash<::kumo::pollux::exec::SpillPartitionId> {
        size_t operator()(const ::kumo::pollux::exec::SpillPartitionId &id) const {
            return kumo::pollux::bits::hashMix(
                id.partitionBitOffset(), id.partitionNumber());
        }
    };
} // namespace std

template<>
struct fmt::formatter<kumo::pollux::exec::SpillPartitionId>
        : formatter<std::string> {
    auto format(kumo::pollux::exec::SpillPartitionId s, format_context &ctx)
    const {
        return formatter<std::string>::format(s.toString(), ctx);
    }
};
